import sensor,image,lcd,time
import KPU as kpu
from machine import UART
from fpioa_manager import fm
from Maix import GPIO

lcd.init(freq=15000000)
#lcd.rotation(1)
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.set_vflip(1)    #设置摄像头后置
sensor.run(1)

task = kpu.load("/sd/yolov3.kmodel")


anchor = (0.3691, 0.2575, 0.4528, 0.4739, 0.6996, 0.5839, 0.9164, 0.9705, 2.3887, 1.7223)
sensor.set_windowing((224, 224))
a = kpu.init_yolo2(task, 0.5, 0.3, 5, anchor)
classes=["9","1","4","2","3","8","5","6","7" ]

target_id = 0
start_flag = 0
pause_flag = 1
recognition_flag = 0

fm.register(22,fm.fpioa.GPIO0)
fm.register(23,fm.fpioa.GPIO1)
fm.register(24,fm.fpioa.GPIO2)
fm.register(25,fm.fpioa.GPIO3)

key0=GPIO(GPIO.GPIO0,GPIO.IN, GPIO.PULL_DOWN)
key1=GPIO(GPIO.GPIO1,GPIO.IN, GPIO.PULL_DOWN)
left_out=GPIO(GPIO.GPIO2,GPIO.OUT)
right_out=GPIO(GPIO.GPIO3,GPIO.OUT)

key0_Press = 0
key1_Press = 0
lastKey0_State = 1
lastKey1_State = 1
def key_scan():
    global key0_Press
    global key1_Press
    global lastKey0_State
    global lastKey1_State
    
    val0 = key0.value()
    val1 = key1.value()

    if lastKey0_State == 0 and val0 == 1:
        key0_Press = 1
    else:
        key0_Press = 0
    lastKey0_State = val1

    if lastKey1_State == 0 and val1 == 1:
        key1_Press += 1
    elif key1_Press >= 2:
        key1_Press = 0
    lastKey1_State = val1


def start_statu_pro():
    global target_id
    for x in range(0,30):
        img = sensor.snapshot()
        code = kpu.run_yolo2(task, img)

        if code:
            for i in code:
                a=img.draw_rectangle(i.rect())
                a = lcd.display(img)
                list1=list(i.rect())
                b=(list1[0]+list1[2])/2
                c=(list1[1]+list1[3])/2
                #print("物体是：",classes[i.classid()])
                #print("概率为：",100.00*i.value())
                #print("坐标为：",b,c)
                for i in code:
                    lcd.draw_string(0, 0, "Input...", lcd.RED, lcd.WHITE)
                    lcd.draw_string(i.x(), i.y(), classes[i.classid()], lcd.RED, lcd.WHITE)
                    lcd.draw_string(i.x(), i.y()+12, '%f'%i.value(), lcd.RED, lcd.WHITE)
                    lcd.draw_string(50, 200,str(b), lcd.RED, lcd.WHITE)
                    lcd.draw_string(110, 200,str(c), lcd.RED, lcd.WHITE)

            code = sorted(code, key=(lambda x:x.value()), reverse=True)
            target_id = code[0].classid()
            if classes[target_id] == 1:
                left_out.value(1)
                right_out.value(0)
            elif classes[target_id] == 2:
                left_out.value(1) 
                right_out.value(0)
        else:
            a = lcd.display(img)

def pause_statu_pro():
    while(True):
        img = sensor.snapshot()
        code = kpu.run_yolo2(task, img)
        if code:
            for i in code:
                a=img.draw_rectangle(i.rect())
                a = lcd.display(img)
                list1=list(i.rect())
                b=(list1[0]+list1[2])/2
                c=(list1[1]+list1[3])/2
                #print("物体是：",classes[i.classid()])
                #print("概率为：",100.00*i.value())
                #print("坐标为：",b,c)
                for i in code:
                    lcd.draw_string(0, 0, "Pause...", lcd.RED, lcd.WHITE)
                    lcd.draw_string(i.x(), i.y(), classes[i.classid()], lcd.RED, lcd.WHITE)
                    lcd.draw_string(i.x(), i.y()+12, '%f'%i.value(), lcd.RED, lcd.WHITE)
                    lcd.draw_string(50, 200,str(b), lcd.RED, lcd.WHITE)
                    lcd.draw_string(110, 200,str(c), lcd.RED, lcd.WHITE)

        else:
         a = lcd.display(img)

def recognition_statu_pro(direct):
    global target_id
    for x in range(0,5):
        img = sensor.snapshot()
        code = kpu.run_yolo2(task, img)
        if code:
            for i in code:
                a=img.draw_rectangle(i.rect())
                a = lcd.display(img)
                list1=list(i.rect())
                b=(list1[0]+list1[2])/2
                c=(list1[1]+list1[3])/2
                #print("物体是：",classes[i.classid()])
                #print("概率为：",100.00*i.value())
                #print("坐标为：",b,c)
                for i in code:
                    lcd.draw_string(i.x(), i.y(), classes[i.classid()], lcd.RED, lcd.WHITE)
                    lcd.draw_string(i.x(), i.y()+12, '%f'%i.value(), lcd.RED, lcd.WHITE)
                    lcd.draw_string(50, 200,str(b), lcd.RED, lcd.WHITE)
                    lcd.draw_string(110, 200,str(c), lcd.RED, lcd.WHITE)
                if i.classid() == target_id:
                    if direct == 1:
                        left_out.value(1)
                        right_out.value(0)
                    else:
                        left_out.value(0)
                        right_out.value(1)

        else:
         a = lcd.display(img)

def state_transitions():
    global start_flag
    global pause_flag
    global recognition_flag
    if key0_Press:
        start_flag = True
        pause_flag = False
        recognition_flag = False
    elif key1_Press == 1:
        recognition_flag = 1
        start_flag = False
        pause_flag = False
    elif key1_Press == 2:
        recognition_flag = 2
        start_flag = False
        pause_flag = False

def main():
    global start_flag
    global pause_flag
    global recognition_flag

    while (True):
        key_scan()
        state_transitions()
        if start_flag:
            start_statu_pro()
        elif pause_flag:
            pause_statu_pro()
        elif recognition_flag:
            recognition_statu_pro(recognition_flag)

if __name__ == "__main__":
    main()

